import type { BamlRuntime, FunctionResult, BamlCtxManager, ClientRegistry, Image, Audio, Pdf, Video, FunctionLog,
HTTPRequest } from "@boundaryml/baml"
import { toBamlError, BamlStream, BamlAbortError, Collector } from "@boundaryml/baml"
import type { Checked, Check, RecursivePartialNull as MovedRecursivePartialNull } from "./types"
import type { partial_types } from "./partial_types"
import type * as types from "./types"
import type {
{%- for t in types %}{{ t }}{% if !loop.last %}, {% endif %}{% endfor -%}
} from "./types"
import type TypeBuilder from "./type_builder"
import { AsyncHttpRequest, AsyncHttpStreamRequest } from "./async_request"
import { LlmResponseParser, LlmStreamParser } from "./parser"
import { DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX,
DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME } from "./globals"
import type * as events from "./events"

/**
* @deprecated Use RecursivePartialNull from 'baml_client/types' instead.
*/
export type RecursivePartialNull<T> = MovedRecursivePartialNull<T>

    type TickReason = "Unknown";

    type BamlCallOptions<WatchersT = never> = {
    tb?: TypeBuilder
    clientRegistry?: ClientRegistry
    collector?: Collector | Collector[]
    env?: Record<string, string | undefined>
      tags?: Record<string, string>
        signal?: AbortSignal
        onTick?: (reason: TickReason, log: FunctionLog | null) => void
        watchers?: WatchersT
        }

        export class BamlAsyncClient {
        private runtime: BamlRuntime
        private ctxManager: BamlCtxManager
        private streamClient: BamlStreamClient
        private httpRequest: AsyncHttpRequest
        private httpStreamRequest: AsyncHttpStreamRequest
        private llmResponseParser: LlmResponseParser
        private llmStreamParser: LlmStreamParser
        private bamlOptions: BamlCallOptions

        constructor(runtime: BamlRuntime, ctxManager: BamlCtxManager, bamlOptions?: BamlCallOptions) {
        this.runtime = runtime
        this.ctxManager = ctxManager
        this.streamClient = new BamlStreamClient(runtime, ctxManager, bamlOptions)
        this.httpRequest = new AsyncHttpRequest(runtime, ctxManager)
        this.httpStreamRequest = new AsyncHttpStreamRequest(runtime, ctxManager)
        this.llmResponseParser = new LlmResponseParser(runtime, ctxManager)
        this.llmStreamParser = new LlmStreamParser(runtime, ctxManager)
        this.bamlOptions = bamlOptions || {}
        }

        withOptions(bamlOptions: BamlCallOptions) {
        return new BamlAsyncClient(this.runtime, this.ctxManager, bamlOptions)
        }

        get stream() {
        return this.streamClient
        }

        get request() {
        return this.httpRequest
        }

        get streamRequest() {
        return this.httpStreamRequest
        }

        get parse() {
        return this.llmResponseParser
        }

        get parseStream() {
        return this.llmStreamParser
        }

        {% for func in functions %}
        async {{ func.name }}(
        {% for (name, arg_type) in func.args -%}
        {{name}}{% if arg_type.meta().is_optional() %}?{% endif %}: {{arg_type.serialize_type(pkg)}},
        {%- endfor %}
        __baml_options__?: BamlCallOptions<{{ func.event_type_param() }}>
        ): Promise<{{func.return_type.serialize_type(pkg)}}> {
          try {
          const options = { ...this.bamlOptions, ...(__baml_options__ || {}) }
          const signal = options.signal;

          if (signal?.aborted) {
          throw new BamlAbortError('Operation was aborted', signal.reason);
          }

          // Check if onTick is provided - route through streaming if so
          if (options.onTick) {
          const stream = this.stream.{{ func.name }}(
          {% for (name, arg_type) in func.args -%}
          {{name}},
          {%- endfor %}
          __baml_options__
          );

          return await stream.getFinalResponse();
          }

          const collector = options.collector ? (Array.isArray(options.collector) ? options.collector :
          [options.collector]) : [];
          const rawEnv = __baml_options__?.env ? { ...process.env, ...__baml_options__.env } : { ...process.env };
          const env: Record<string, string> = Object.fromEntries(
            Object.entries(rawEnv).filter(([_, value]) => value !== undefined) as [string, string][]
            );
            const raw = await this.runtime.callFunction(
            "{{func.name}}",
            {
            {% for (name, arg_type) in func.args -%}
            "{{name}}": {{name}}{% if arg_type.meta().is_optional() %}?? null{% endif %}{% if !loop.last %},{% endif %}
            {%- endfor %}
            },
            this.ctxManager.cloneContext(),
            options.tb?.__tb(),
            options.clientRegistry,
            collector,
            options.tags || {},
            env,
            signal,
            options.watchers,
            )
            return raw.parsed(false) as {{func.return_type.serialize_type(pkg)}}
            } catch (error) {
            throw toBamlError(error);
            }
            }
            {% endfor %}
            }

            class BamlStreamClient {
            private runtime: BamlRuntime
            private ctxManager: BamlCtxManager
            private bamlOptions: BamlCallOptions

            constructor(runtime: BamlRuntime, ctxManager: BamlCtxManager, bamlOptions?: BamlCallOptions) {
            this.runtime = runtime
            this.ctxManager = ctxManager
            this.bamlOptions = bamlOptions || {}
            }

            {% for func in functions %}
            {{ func.name }}(
            {% for (name, arg_type) in func.args -%}
            {{name}}{% if arg_type.meta().is_optional() %}?{% endif %}: {{arg_type.serialize_type(pkg)}},
            {%- endfor %}
            __baml_options__?: BamlCallOptions<{{ func.event_type_param() }}>
            ): BamlStream<{{ func.stream_return_type.serialize_type(pkg) }}, {{ func.return_type.serialize_type(pkg) }}>
              {
              try {
              const options = { ...this.bamlOptions, ...(__baml_options__ || {}) }
              const signal = options.signal;

              if (signal?.aborted) {
              throw new BamlAbortError('Operation was aborted', signal.reason);
              }

              let collector = options.collector ? (Array.isArray(options.collector) ? options.collector :
              [options.collector]) : [];

              let onTickWrapper: (() => void) | undefined;

              // Create collector and wrap onTick if provided
              if (options.onTick) {
              const tickCollector = new Collector("on-tick-collector");
              collector = [...collector, tickCollector];

              onTickWrapper = () => {
              const log = tickCollector.last;
              if (log) {
              try {
              options.onTick!("Unknown", log);
              } catch (error) {
              console.error("Error in onTick callback for {{func.name}}", error);
              }
              }
              };
              }

              const rawEnv = __baml_options__?.env ? { ...process.env, ...__baml_options__.env } : { ...process.env };
              const env: Record<string, string> = Object.fromEntries(
                Object.entries(rawEnv).filter(([_, value]) => value !== undefined) as [string, string][]
                );
                const raw = this.runtime.streamFunction(
                "{{func.name}}",
                {
                {% for (name, arg_type) in func.args -%}
                "{{name}}": {{name}}{% if arg_type.meta().is_optional() %} ?? null{% endif %}{% if !loop.last %},{%
                endif %}
                {%- endfor %}
                },
                undefined,
                this.ctxManager.cloneContext(),
                options.tb?.__tb(),
                options.clientRegistry,
                collector,
                options.tags || {},
                env,
                signal,
                onTickWrapper,
                )
                return new BamlStream<{{ func.stream_return_type.serialize_type(pkg) }}, {{
                  func.return_type.serialize_type(pkg) }}>(
                  raw,
                  (a): {{ func.stream_return_type.serialize_type(pkg) }} => a,
                  (a): {{ func.return_type.serialize_type(pkg) }} => a,
                  this.ctxManager.cloneContext(),
                  options.signal,
                  )
                  } catch (error) {
                  throw toBamlError(error);
                  }
                  }
                  {% endfor %}
                  }

                  export const b = new BamlAsyncClient(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME,
                  DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX)
